Cache launch_configuration per kernel - #605
Open
michel2323 wants to merge 1 commit into
Open
Conversation
`launch_configuration` issues a `zeKernelGetProperties` round-trip, and KernelAbstractions calls it on every dispatch of a kernel whose workgroupsize is `DynamicSize` and unspecified at the call site (src/oneAPIKernels.jl). The query therefore lands on the per-launch path, where it dominates: on a Data Center GPU Max 1550, a no-op kernel costs 96 us per launch as-is and 9.7 us with the workgroupsize pinned, a 9.9x difference that is entirely host-side. Launch-bound workloads pay it on every kernel they dispatch. The returned group size depends only on the kernel and its device. A ZeKernel's handle and module are fixed for its lifetime, and maxGroupSize / maxTotalGroupSize are static properties of the compiled kernel, so neither the ndrange nor any argument value can change the result. Keyed weakly on the ZeKernel rather than on its raw handle: a handle is unique only among live kernels, so a handle-keyed entry could be inherited by a later kernel that reuses a destroyed kernel's address. A weak key ties the entry to the kernel's lifetime and needs no hook in oneL0's finalizer. Measured with the cache in place: 96 us -> 62 us per launch. The remainder is the re-partition and context rebuild that follow the query in the KernelAbstractions backend, which this does not address.
Contributor
|
Your PR requires formatting changes to meet the project's style guidelines. Click here to view the suggested changes.diff --git a/src/compiler/execution.jl b/src/compiler/execution.jl
index c3a7429..8df7b33 100644
--- a/src/compiler/execution.jl
+++ b/src/compiler/execution.jl
@@ -228,7 +228,7 @@ end
# unique among *live* kernels, so a handle key could let a destroyed kernel's
# entry be inherited by a later kernel that reuses its address. The weak key ties
# the entry to the kernel's own lifetime and needs no hook in oneL0's finalizer.
-const _launch_config_cache = WeakKeyDict{ZeKernel,Int}()
+const _launch_config_cache = WeakKeyDict{ZeKernel, Int}()
function launch_configuration(kernel::HostKernel{F,TT}) where {F,TT}
fun = kernel.fun
@@ -239,7 +239,7 @@ function launch_configuration(kernel::HostKernel{F,TT}) where {F,TT}
return config
end
-function _launch_configuration_uncached(kernel::HostKernel{F,TT}) where {F,TT}
+function _launch_configuration_uncached(kernel::HostKernel{F, TT}) where {F, TT}
# Level Zero's zeKernelSuggestGroupSize provides a launch configuration
# that exactly cover the input size. This can result in very awkward
# configurations, so roll our own version that behaves like CUDA's |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
launch_configurationissues azeKernelGetPropertiesround-trip. KernelAbstractions calls it on every dispatch of a kernel whose workgroupsize isDynamicSizeand unspecified at the call site (src/oneAPIKernels.jl), so the query lands on the per-launch path rather than a setup path, and there it dominates.Measured on a Data Center GPU Max 1550 with a no-op kernel (
a[i] += 1f0, n=1024, best of 5 batches of 2000 launches):workgroupsizepassed at call siteStaticSize@oneapi, no KAA 9.9x spread, entirely host-side — the GPU does the same trivial work in every row. Launch-bound workloads pay this on every kernel they dispatch.
The change
The returned group size depends only on the kernel and its device. A
ZeKernel's handle and module are fixed for its lifetime, andmaxGroupSize/maxTotalGroupSizeare static properties of the compiled kernel, so neither the ndrange nor any argument value can change the result. It is computed once per kernel and remembered.Keyed weakly on the
ZeKernelrather than on its raw handle: a handle is unique only among live kernels, so a handle-keyed entry could be inherited by a later kernel that reuses a destroyed kernel's address. A weak key ties the entry to the kernel's lifetime and needs no hook inoneL0's finalizer.Effect, and what it does not fix
With the cache in place the same benchmark goes 9.63e-5 s -> 6.25e-5 s per launch, and the cache is confirmed to hold a single entry after 501 launches.
That is about a third of the gap. The remaining ~53 us is the re-partition and context rebuild that follow the query in the KA backend (
src/oneAPIKernels.jl, theKA.partition/KA.mkcontextpair after thelaunch_configurationcall), which this PR deliberately does not touch. Passingworkgroupsizeat the call site skips that whole branch and reaches 1.99e-5 s, so there is more to win there separately.For reference, CUDA.jl's KA backend also calls
launch_configurationper launch; the difference is that its occupancy API is cheap wherezeKernelGetPropertiesis not. Caching here seemed preferable to changing every caller.